package yycore
import chisel3._
import chisel3.util._
import common.Constants.{XLEN, fuOpType}
import utils.SignExt

object MDUOpType {
  def mul    = "b0000".U
  def mulh   = "b0001".U
  def mulhsu = "b0010".U
  def mulhu  = "b0011".U
  def div    = "b0100".U
  def divu   = "b0101".U
  def rem    = "b0110".U
  def remu   = "b0111".U

  def mulw   = "b1000".U
  def divw   = "b1100".U
  def divuw  = "b1101".U
  def remw   = "b1110".U
  def remuw  = "b1111".U

  def isDiv(op: UInt) = op(2)
  def isDivSign(op: UInt) = isDiv(op) && !op(0)
  def isW(op: UInt) = op(3)
}

class OpDataReqIO(val len: Int) extends Bundle {
  val op1 = Output(UInt(len.W))   // 被除数
  val op2 = Output(UInt(len.W))  // 除数
}
class OpDataRespIO(val len: Int) extends Bundle {
  val dataout = Output(UInt(len.W))
}
// 64位多周期除法器
// class DIVModule(len: Int = 64) extends Module {
//   val io = IO(new Bundle(){
//     val in = Flipped(Decoupled(new DivReqIO(len)))
//     val out = Decoupled(new DivRespIO(len))
//     val sign = Input(Bool())
//   })

//   val idle :: busy :: out :: Nil = Enum(3)
//   val reg_status = RegInit(idle)
//   val newReq = (reg_status === idle) && io.in.fire()


//   val cntReg = RegInit(0.U(7.W))
//   val reg_divisor = RegInit(0.U(128.W))
//   val reg_remainder = RegInit(0.U(128.W))
//   //val reg_quotient = RegInit(0.U(DataBits.W))


//   // 负数补码转换为原码
//   def abs(a: UInt, sign: Bool): (Bool, UInt) = {
//     val s = a(len - 1) && sign
//     (s, Mux(s, -a, a))
//   }

//   val (a, b) = (io.in.bits.op1, io.in.bits.op2)
//   val divBy0 = b === 0.U(len.W)
//   val (aSign, aVal) = abs(a, io.sign)
//   val (bSign, bVal) = abs(b, io.sign)
//   val aSignReg = RegEnable(aSign, newReq)
//   val qSignReg = RegEnable((aSign ^ bSign) && !divBy0, newReq)
//   //printf("aVal=%x, newReq=%x,reg_status=%x\n", aVal, newReq,reg_status)
//   // 初始化
//   when(newReq){
//     reg_status := busy
//     cntReg := 64.U
//     //printf("aVal=%x\n", aVal)
//     reg_remainder := Cat(Fill(64, 0.U), aVal)
//     reg_divisor := Cat(bVal, Fill(64, 0.U))
//     //reg_quotient := 0.U
//   }


//   when(reg_status === busy && cntReg =/= 0.U){
//     cntReg := cntReg - 1.U
//     val shiftl1 = reg_remainder << 1
//     val compare = shiftl1 > reg_divisor
//     when(compare){
//       reg_remainder := shiftl1 - reg_divisor + 1.U
//     }.otherwise{
//       reg_remainder := shiftl1
//     }

//   }.elsewhen(reg_status === busy && cntReg === 0.U){
//     reg_status := out
//   }

//   io.in.ready := reg_status === idle
//   io.out.valid := Mux(reg_status === out, true.B, false.B)
//   val resQ = Mux(qSignReg, -reg_remainder(63, 0), reg_remainder(63, 0))
//   val resR = Mux(aSignReg, -reg_remainder(127, 64), reg_remainder(127, 64))
//   io.out.bits.div_out := resQ

//   when(reg_status === out && io.out.fire()){
//     //printf("OUT MODE\n")
//     reg_status := idle
//   }
// }

class DIVModule(len: Int = 64) extends Module {
  val io = IO(new Bundle(){
    val in = Flipped(Decoupled(new OpDataReqIO(len)))
    val out = Decoupled(new OpDataRespIO(len))
    val sign = Input(Bool())
  })

  def abs(a: UInt, sign: Bool): (Bool, UInt) = {
    val s = a(len - 1) && sign
    (s, Mux(s, -a, a))
  }

  val s_idle :: s_log2 :: s_shift :: s_compute :: s_finish :: Nil = Enum(5)
  val state = RegInit(s_idle)
  val newReq = (state === s_idle) && io.in.fire()

  val (a, b) = (io.in.bits.op1, io.in.bits.op2)
  val divBy0 = b === 0.U(len.W)

  val shiftReg = Reg(UInt((1 + len * 2).W))
  val hi = shiftReg(len * 2, len)
  val lo = shiftReg(len - 1, 0)

  val (aSign, aVal) = abs(a, io.sign)
  val (bSign, bVal) = abs(b, io.sign)
  val aSignReg = RegEnable(aSign, newReq)
  val qSignReg = RegEnable((aSign ^ bSign) && !divBy0, newReq)
  val bReg = RegEnable(bVal, newReq)
  val aValx2Reg = RegEnable(Cat(aVal, "b0".U), newReq)

  val cnt = Counter(len)
  when (newReq) {
    state := s_log2
  } .elsewhen (state === s_log2) {
    val canSkipShift = (len.U | Log2(bReg)) - Log2(aValx2Reg)
    cnt.value := Mux(divBy0, 0.U, Mux(canSkipShift >= (len-1).U, (len-1).U, canSkipShift))
    state := s_shift
  } .elsewhen (state === s_shift) {
    shiftReg := aValx2Reg << cnt.value
    state := s_compute
  } .elsewhen (state === s_compute) {
    val enough = hi.asUInt >= bReg.asUInt
    shiftReg := Cat(Mux(enough, hi - bReg, hi)(len - 1, 0), lo, enough)
    cnt.inc()
    when (cnt.value === (len-1).U) { state := s_finish }
  } .elsewhen (state === s_finish) {
    state := s_idle
  }

  val r = hi(len, 1)
  val resQ = Mux(qSignReg, -lo, lo)
  val resR = Mux(aSignReg, -r, r)
  io.out.bits.dataout := Cat(resR, resQ)

  io.out.valid := state === s_finish
  io.in.ready := (state === s_idle)
}

class MulModule(len: Int = 64) extends Module {
  val io = IO(new Bundle() {
    val in = Flipped(Decoupled(new OpDataReqIO(len)))
    val out = Decoupled(new OpDataRespIO(2*len))
  })

  val (src1,  src2) = (io.in.bits.op1, io.in.bits.op2)

  val mulOut = (src1 * src2).asUInt()

  io.in.ready := true.B
  io.out.valid := io.in.valid
  io.out.bits.dataout := mulOut
}

class MDU extends Module {
  val io = IO(new Bundle() {
    val in = Flipped(Decoupled(new Bundle {
      val src1 = Output(UInt(XLEN.W))
      val src2 = Output(UInt(XLEN.W))
      val fun = Output(UInt(fuOpType.W))
    }))
    val out = Decoupled(Output(UInt(XLEN.W)))
  })

  val (valid, src1, src2, fun) = (io.in.valid, io.in.bits.src1, io.in.bits.src2, io.in.bits.fun)
  def access(valid: Bool, src1: UInt, src2: UInt, fun: UInt): UInt = {
    this.valid := valid
    this.src1 := src1
    this.src2 := src2
    this.fun := fun
    io.out.bits
  }

  val isDiv = MDUOpType.isDiv(fun)
  val isW = MDUOpType.isW(fun)

  //val mul = Module(new MulModule(XLEN))
  val mul = Module(new BoothMulModule(XLEN))
  val div = Module(new DIVModule(XLEN))
  List(mul.io, div.io).map { case x =>
    x.out.ready := io.out.ready
  }

  /**
   * DIV
   */
  div.io.in.bits.op1 := MuxCase(0.U, Array(
    (fun === MDUOpType.div || fun === MDUOpType.divu)    -> src1,
    (fun === MDUOpType.divw)   -> Cat(Fill(32, src1(31)), src1(31, 0)),
    (fun === MDUOpType.divuw)  -> Cat(Fill(32, 0.U), src1(31, 0)),
  ))
  div.io.in.bits.op2 := MuxCase(0.U, Array(
    (fun === MDUOpType.div || fun === MDUOpType.divu)    -> src2,
    (fun === MDUOpType.divw)   -> Cat(Fill(32, src2(31)), src2(31, 0)),
    (fun === MDUOpType.divuw)  -> Cat(Fill(32, 0.U), src2(31, 0)),
  ))
  div.io.sign := MuxCase(false.B, Array(
    (fun === MDUOpType.div  || fun === MDUOpType.divw)    -> true.B,
    (fun === MDUOpType.divu || fun === MDUOpType.divuw)   -> false.B,
  ))

  /**
   * MUL
   */
  mul.io.in.bits.op1 := Mux(!isDiv, src1, 0.U)
  mul.io.in.bits.op2 := Mux(!isDiv, src2, 0.U)

  mul.io.in.valid := io.in.valid && !isDiv
  div.io.in.valid := io.in.valid && isDiv

  val mulRes = Mux(fun(1,0) === "b00".U, mul.io.out.bits.dataout(XLEN-1,0), mul.io.out.bits.dataout(2*XLEN-1,XLEN))
  val divRes = div.io.out.bits.dataout

  val res = Mux(isDiv, divRes, mulRes)
  io.out.valid := Mux(isDiv, div.io.out.valid, mul.io.out.valid)
  io.out.bits := Mux(isW, SignExt(res(31,0), XLEN), res)
  io.in.ready := mul.io.in.ready && div.io.in.ready
}

trait BoothConst {
  val neg = Bool()
  // base 4 idx booth
  def booth_enc(code: UInt): UInt = {
    val re = VecInit((0 until 4).map(i => false.B)) // neg, two, zero, one
    re(0) := code(2)
    when(code === "b011".U || code === "b100".U) { re(1) := true.B }
      .elsewhen(code === "b000".U || code === "b000".U) { re(2) := true.B }
      .otherwise { re(3) := true.B }
    re.asUInt()
  }
  def gen_prod(multiplicand: UInt, enc: UInt) : UInt = {
    //val enc = VecInit((0 until 4))
    val prod_pre = Wire(UInt((2*XLEN).W))
    when(enc(1)) { prod_pre := Cat(multiplicand(XLEN-1), multiplicand, "b0".U) }
    .elsewhen(enc(2)) { prod_pre := 0.asUInt(XLEN.W) }
      .otherwise { prod_pre := multiplicand }
    val prod = Mux(enc(0), ((~prod_pre).asUInt() + "b1".U), prod_pre)
    prod
  }
}

class BoothMulModule(len: Int) extends Module with BoothConst {
  val io = IO(new Bundle() {
    val in = Flipped(Decoupled(new OpDataReqIO(len)))
    val out = Decoupled(new OpDataRespIO(2*len))
  })

  val (src1,  src2) = (io.in.bits.op1, io.in.bits.op2)



  val s_idle :: s_busy :: s_finish :: Nil = Enum(3)
  val state = RegInit(s_idle)
  val cnt = RegInit(0.asUInt(6.W))
  io.in.ready := state === s_idle
  io.out.valid := state === s_finish
  //
  val multiplicand = RegEnable(next = Cat(src1(len-1), src1), init = 0.asUInt((len+1).W), enable = io.in.fire())
  val multiplier = RegEnable(next = Cat(src1, "b0".U), init = 0.asUInt((len+1).W), enable = io.in.fire())
  //val P = RegInit(0.asUInt((2*len+3).W))
  val Q = RegInit(0.asUInt((2*len).W))
  switch(state){
    is(s_idle){
      when(io.in.fire()) {
        state := s_busy
        cnt := 0.U
        Q := 0.U
        //P := Cat(Fill(len, src2(XLEN-1)), src2, "b0".U)
      }
    }
    is(s_busy) {
      when(cnt === 32.U) {
        state := s_finish
      }
    }
    is(s_finish) {
      when(io.out.fire()) {
        state := s_idle
      }

    }
  }

  when (state === s_busy) {
    val code = multiplier(64, 62)
    val prod = gen_prod(multiplicand, booth_enc(code))
    // left shift
    val nextQ = Cat(Q(2*XLEN - 1, 2), Fill(2, "b0".U)) + prod

    multiplier := Cat(multiplier(XLEN - 2, 0), Fill(2, "b0".U))
    Q := nextQ
    cnt := cnt + 1.U
  }

  io.out.bits.dataout := Q

}